'''
This script produces a map of crystalline regions from a 4D STEM dataset in DigitalMicrograph
To run, first open the 4D dataset in DigitalMicrograph, then place a picker-tool on the 4D dataset so 
a diffration pattern is shown. Then place a circular ROI around the region in the center that should be 
ignored when finding the brightest pixel in each diffraction pattern. 

The ROI must be circular, not oval, or an error will be returned.

This requires a module FFTArrayAnalysis, which requires scipy, skimage, matplotlib, and tqdm 
The module is part of a package which can be installed with pip:
    pip install BenMillerScripts
Code written by Ben Miller. Last Updated July 2024
'''
from benmillerscripts import FFTArrayAnalysis as FAA
import numpy as np

#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
#User-Set Parameters

FAA.maskC_width = 0 #Pixels to mask on the center vertical and horizontal
map_var = 'theta'   #Variable to Map in the colormap ('theta' or 'radius')
                    #theta maps the angle of the brightest FFT spot
                    #radius maps the spacing of the brightest FFT spot
scale = 1           #Optional Scaling factor (should be a positive integer)
                    #This scales the data up in real space, but down in diffraction space

FAA.pre_filter = False       #Optionally filter the 4DSTEM dataset with a median filter (recommended)
showcube = True             #Optionally show the datacube that is analyzed (after cropping, scaling, and masking)
show_grey_maps = False       #Optionally show the 2 maps that are used to generate the color map
subtract_mean_bg = False    #Optionally subtract the mean diffraction pattern from every pattern
draw_color_scale = True    #Optionally produce a circular color scale
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

def get_pickerROI_num():
    '''
    Function to get a picker-tool index number from the front-most data in GMS
    returns: ROI_num    integer index for an region of interest (ROI) in a GMS image
    '''
    #DM Script Embedded in this Python Code
    #The script gets the picker tool number from the front-most image
    # and puts it in a global tag that Python can access
    dm='imagedocument imdoc = GetFrontImageDocument()' +'\n'
    dm+='imageDisplay disp = imdoc.ImageDocumentGetImageModeDisplay()' +'\n'
    dm+='number nR = disp.ImageDisplayCountROIs()' +'\n'
    dm+='if(nR==0) Result("No Picker Tool Present on 4D STEM Dataset")' +'\n'
    dm+='if(nR>1) Result("Warning, Multiple Picker Tools Present on 4D STEM Dataset, Using Number "+disp.ImageDisplayGetROI(0).ROIGetLabel())' +'\n'
    dm+='number ROI_num = disp.ImageDisplayGetROI(0).ROIGetLabel().val()' +'\n'
    dm+='GetPersistentTagGroup().TagGroupSetTagAsLong("Python_temp:out:ROInum",ROI_num)' +'\n'
    #Run DM Script
    try: DM.ExecuteScriptString(dm)
    except: print("No Picker Tool Found on Front-Most Image"); exit()
    #Get # from Tags
    (b,ROI_num) = DM.GetPersistentTagGroup().GetTagAsUInt32('Python_temp:out:ROInum')
    return ROI_num
    
def get_picker_data_and_ROI(image):
    '''
    Function to get a diffraction pattern from the linked picker-tool data and a rectangular ROI within it
    Accepts: image                         the original DM image containing a 4DSTEM dataset- this should have a picker-tool placed on it, be front-most in GMS,
    returns: (top,left,bottom,right)    the coordinates (in pixels) of the rectangular ROI position
             diff_data                    numpy array with the picker tool data (a diffraction pattern)
    '''
    #make sure this image is front-most
    image.ShowImage() 
    #check that this is a 4D dataset
    if image.GetNumDimensions() != 4: DM.OkDialog("Front-Most Image is not a 4D Dataset...\n\nAborting Script") ; exit()
    #Get picker tool diffraction data
    Diff_Name = "("+str(get_pickerROI_num())+") Diffraction of "+image.GetName()
    diff_img = DM.FindImageByName( Diff_Name )
    diff_data = diff_img.GetNumArray()
    sy,sx = diff_data.shape
    if ( diff_img is None ):
        print( 'No image "', Diff_Name, '" was found' )
    #Get ROI coordinates
    imageDisplay = diff_img.GetImageDisplay(0) 
    roi = imageDisplay.GetROI(0)
    try: 
        if roi.IsRectangle:
            top,left,bottom,right = np.array(roi.GetRectangle(),dtype='int32')
            x = (right+left)//2
            y = (bottom+top)//2
            r = (right-left)/2
        if roi.IsCircle():
            x,y,r = np.array(roi.GetCircle(),dtype='int32')
    except: print("Problem Reading ROI on the Picker Tool Image... Script Aborted. Please Place a Circular ROI around the Central Spot"); exit()
    del diff_img, image
    return (x,y,r),diff_data
    
 
    
(x,y,r), picker_diff= get_picker_data_and_ROI(DM.GetFrontImage())
sx,sy = picker_diff.shape
center = (x,y)
FAA.maskP = r*2/sx*100

#Get Data from GMS
image_4D = DM.GetFrontImage()
raw_datacube = image_4D.GetNumArray()
name = image_4D.GetName()
origin, x_scale, scale_unit =  image_4D.GetDimensionCalibration(2, 0)

#Optionally Scale Data (up in real space, down in diff space)
if scale > 1:
    data_s = FAA.STEMx_Crop_to_Center_Square(raw_datacube, center)
    datasquare = np.copy(FAA.Upscale_4D_STEM(data_s, scale))
else: 
    datasquare = np.copy(FAA.STEMx_Crop_to_Center_Square(raw_datacube, center))

#Optionally Subtract Mean Pattern from Data
if subtract_mean_bg: 
    avg_diffraction = np.mean(datasquare,axis=(0,1))
    data_BGSub = datasquare-avg_diffraction
    data_BGSub[data_BGSub<0] = 0
    data_input = data_BGSub
else: 
    data_input = datasquare

#Process Dataset
(RGB_im,direction_image,intensity_image,spacing_image,diffractogram_max) = FAA.STEMx_Process_Cube(image_4D,map_var, im_data = data_input, show_cube = showcube)
DM_RGB = FAA.ShowRGB(RGB_im, name)

#Optionally Show More Maps
if show_grey_maps:
    d_im = DM.CreateImage(np.copy(direction_image))
    d_im.SetName("FFT Spot Direction")
    d_im.SetDimensionCalibration(0,origin, x_scale, scale_unit,0)
    d_im.SetDimensionCalibration(1,origin, x_scale, scale_unit,0)
    d_im.ShowImage()
    del d_im
    d_im = DM.CreateImage(np.copy(intensity_image))
    d_im.SetDimensionCalibration(0,origin, x_scale, scale_unit,0)
    d_im.SetDimensionCalibration(1,origin, x_scale, scale_unit,0)
    d_im.SetName("FFT Spot Intensity")
    d_im.ShowImage()
    del d_im
del image_4D

#Optionally Produce Circular Color Scale
if draw_color_scale: FAA.DrawColorScale()